package geppetto.phraseHMM;

import geppetto.cat.common.StaticTools;

/**
 * Represents a treelis a table that is Number of states * Number of positions
 * 
 * @author javg
 * 
 */
public class Trellis {

	public int _nrStates;

	public int _nrPositions;

	boolean _withNull;

	// This representes the _nrStates vs _nrPositions matrix
	// _ 0 _ 1 _ 2 _ 3 __ 4 _ 5
	// 0 | 0 3 6 9 12 15 | nr_states*(pos)+stateIndex
	// 1 | 1 4 7 10 13 16 |
	// 2 | 2 5 8 11 14 17 |
	// ___________________________
	public double[] _trellis;

	public Trellis(int nrPositions, int nrStates) {
		_trellis = new double[nrStates * nrPositions];
		_nrStates = nrStates;
		_nrPositions = nrPositions;
		// Used for printing purposes
		_withNull = false;
	}

	public Trellis(int nrPositions, int nrStates, boolean withNull) {
		_trellis = new double[nrStates * nrPositions];
		_nrStates = nrStates;
		_nrPositions = nrPositions;
		// Used for printing purposes
		_withNull = withNull;
	}
	
	public Trellis deppCopy(){
		Trellis t = new Trellis(_nrPositions,_nrStates,_withNull);	
		for(int i = 0; i < _trellis.length; i++){
			t._trellis[i] = _trellis[i];
		}
		return t;
	}

	public final void setProb(int position, int state, double prob) {
//		 System.out.println("p" + position + " s " + state);
		_trellis[_nrStates * (position) + state] = prob;
	}

	public final void addProb(int position, int state, double prob) {
		_trellis[_nrStates * (position) + state] += prob;
	}

	public final double getProb(int position, int state) {
		// System.out.println("p" + position + " s " + state);
		return _trellis[_nrStates * (position) + state];
	}

	// Initialize to -1 because i want to have an error if we
	// go out of the arrays
	public final void clear() {
		java.util.Arrays.fill(_trellis, -1);
	}

	public final void clear(double value) {
		java.util.Arrays.fill(_trellis, value);
	}

	public final void clear(int positions, int states) {
		java.util.Arrays.fill(_trellis, 0, states * positions, -1);
	}

	public final void print(int maxPosition, int maxState) {
		int middle = maxState / 2;
		System.out.print("\t");
		for (int pos = 0; pos < maxPosition; pos++) {
			System.out.print("   p" + pos + "   \t");
		}
		System.out.println();
		for (int state = 0; state < maxState; state++) {
			if (_withNull) {
				if (state >= middle)
					System.out.print("ns" + (state - middle) + "\t");
				else
					System.out.print("s" + state + "\t");
			} else {
				System.out.print("s" + state + "\t");
			}
			for (int pos = 0; pos < maxPosition; pos++) {
				System.out.print(StaticTools.prettyPrint(getProb(pos, state),
						"00.00E00", 4)
						+ "\t");
			}
			System.out.println();
		}
	}

	public final void printWithSum(int maxPosition, int maxState) {
		printWithSum("", maxPosition, maxState);
	}
	public final void printWithSum(String name,int maxPosition, int maxState) {
		System.out.println(name);
		int middle = maxState / 2;
		System.out.print("\t");
		double totalSum = 0;
		double[] columnSum = new double[maxPosition];
		for (int pos = 0; pos < maxPosition; pos++) {
			System.out.print("   p" + pos + "   \t");
		}
		System.out.print("   sum");
		System.out.println();
		for (int state = 0; state < maxState; state++) {
			if (_withNull) {
				if (state >= middle)
					System.out.print("ns" + (state - middle) + "\t");
				else
					System.out.print("s" + state + "\t");
			} else {
				System.out.print("s" + state + "\t");
			}
			double rowSum =0;
			for (int pos = 0; pos < maxPosition; pos++) {
				System.out.print(StaticTools.prettyPrint(getProb(pos, state),
						"00.00E00", 4)
						+ "\t");
				columnSum[pos] += getProb(pos, state);
				rowSum += getProb(pos, state);
			}
			totalSum += rowSum;
			System.out.print(StaticTools.prettyPrint(rowSum,
					"00.00E00", 4));
			System.out.println();
		}
		System.out.print("sum\t");
		for (int pos = 0; pos < maxPosition; pos++) {
			System.out.print(StaticTools.prettyPrint(columnSum[pos],
					"00.00E00", 4)+"\t");
		}
		System.out.print(StaticTools.prettyPrint(totalSum,
				"00.00E00", 4));
		System.out.println();
	}
	
	
	
	public final void printAll() {
		for (int state = 0; state < _nrStates; state++) {
			for (int pos = 0; pos < _nrPositions; pos++) {
				System.out.print(getProb(pos, state) + " ");
			}
			System.out.println();
		}
	}

	/**Return a new Trellis of the same size as a and b
	 * which is the entry difference between a and b
	 * @param a
	 * @param b
	 * @return
	 */
	
	public static Trellis difference(Trellis a, Trellis b){
		Trellis diff = new Trellis(a._nrPositions,a._nrStates);
		for(int i = 0; i < diff._trellis.length; i++){
			diff._trellis[i] = a._trellis[i] - b._trellis[i];
		}
		return diff;
	}
	
	public static void main(String[] args) {
		int nrStates = 2;
		int nrPositions = 4;
		Trellis t = new Trellis(4, 2);
		t.print(4, 2);
		for (int sourceIndex = 0; sourceIndex < nrStates; sourceIndex++) {
			t.setProb(nrPositions - 1, sourceIndex, 1);
		}
		System.out.println("After");
		t.print(4, 2);

	}
	
	
}
